"""
    CovarianceLocalization

Advanced covariance localization methods for background error modeling in GSI.
This module implements sophisticated localization functions that taper covariance
matrices to remove spurious long-range correlations while preserving physically
meaningful short-range correlations.

Covariance localization is critical for:
- Ensemble-based data assimilation methods (EnKF, LETKF)
- Hybrid variational-ensemble systems (Hybrid 3D/4D-Var)  
- Sample covariance regularization in limited ensemble systems
- Background error covariance conditioning
- Spurious correlation removal in high-dimensional systems

Mathematical Framework:
Localized covariance matrices are computed as:
```
C_localized = C_raw ⊙ ρ(r/L_loc)
```

Where:
- C_raw is the raw (sample or climatological) covariance
- ρ(r/L_loc) is the localization function
- ⊙ denotes element-wise (Schur) product
- r is the distance between grid points
- L_loc is the localization radius

The module implements multiple localization functions:
- Gaspari-Cohn fifth-order piecewise rational function
- Exponential decay functions for different correlation structures
- Adaptive localization with flow-dependent radius
- Multi-scale localization for hierarchical systems
- Distance-weighted localization with various metrics

Key Features:
- Computationally efficient sparse localization matrices
- Support for multiple distance metrics (Euclidean, great circle, etc.)
- Adaptive localization radius based on ensemble spread
- Multi-variable localization with cross-variable tapering
- Integration with ensemble covariance estimation methods
"""
module CovarianceLocalization

using LinearAlgebra
using SparseArrays
using StaticArrays

export AbstractLocalizationFunction, GaspariCohnFunction, ExponentialFunction
export AdaptiveLocalization, MultiScaleLocalization, DistanceMetric
export LocalizationOperator, LocalizationConfiguration
export apply_localization, compute_localization_matrix, compute_localized_covariance
export create_localization_operator, validate_localization_function

"""
    DistanceMetric

Enumeration of distance metrics for covariance localization.
"""
@enum DistanceMetric begin
    EUCLIDEAN       # Standard Euclidean distance
    GREAT_CIRCLE    # Great circle distance on sphere
    MANHATTAN       # Manhattan (L1) distance  
    MAXIMUM         # Maximum (L∞) distance
    MAHALANOBIS     # Mahalanobis distance with metric tensor
end

"""
    AbstractLocalizationFunction{T<:AbstractFloat}

Abstract base type for localization functions ρ(r).
"""
abstract type AbstractLocalizationFunction{T<:AbstractFloat} end

"""
    GaspariCohnFunction{T<:AbstractFloat} <: AbstractLocalizationFunction{T}

Gaspari-Cohn fifth-order piecewise rational localization function.

The Gaspari-Cohn function provides C² continuity and compact support:
```
ρ(r) = { -1/4(r/c)⁵ + 1/2(r/c)⁴ + 5/8(r/c)³ - 5/3(r/c)² + 1,  0 ≤ r/c ≤ 1
       { 1/12(r/c)⁵ - 1/2(r/c)⁴ + 5/8(r/c)³ + 5/3(r/c)² - 5(r/c) + 4 - 2/3(c/r),  1 < r/c ≤ 2  
       { 0,                                                          r/c > 2

Where c is the localization radius (cutoff distance).

# Fields
- `cutoff_radius::T`: Localization cutoff radius L_loc
- `cutoff_ratio::T`: Ratio for secondary cutoff (default: 2.0)
"""
struct GaspariCohnFunction{T<:AbstractFloat} <: AbstractLocalizationFunction{T}
    cutoff_radius::T
    cutoff_ratio::T
    
    function GaspariCohnFunction{T}(cutoff_radius::T, cutoff_ratio::T = T(2)) where T
        @assert cutoff_radius > 0 "Cutoff radius must be positive"
        @assert cutoff_ratio >= 1 "Cutoff ratio must be ≥ 1"
        new{T}(cutoff_radius, cutoff_ratio)
    end
end

GaspariCohnFunction(cutoff_radius::T, cutoff_ratio::T = T(2)) where T = 
    GaspariCohnFunction{T}(cutoff_radius, cutoff_ratio)

"""
    ExponentialFunction{T<:AbstractFloat} <: AbstractLocalizationFunction{T}

Exponential decay localization function: ρ(r) = exp(-r²/2L²).

# Fields
- `length_scale::T`: Characteristic length scale L
- `decay_rate::T`: Exponential decay rate parameter
- `cutoff_threshold::T`: Cutoff threshold for numerical efficiency
"""
struct ExponentialFunction{T<:AbstractFloat} <: AbstractLocalizationFunction{T}
    length_scale::T
    decay_rate::T
    cutoff_threshold::T
    
    function ExponentialFunction{T}(length_scale::T, decay_rate::T = T(1), 
                                   cutoff_threshold::T = T(1e-6)) where T
        @assert length_scale > 0 "Length scale must be positive"
        @assert decay_rate > 0 "Decay rate must be positive"
        new{T}(length_scale, decay_rate, cutoff_threshold)
    end
end

ExponentialFunction(length_scale::T, args...; kwargs...) where T = 
    ExponentialFunction{T}(length_scale, args...; kwargs...)

"""
    LocalizationConfiguration{T<:AbstractFloat}

Configuration parameters for covariance localization operations.

# Fields
- `base_radius::T`: Base localization radius
- `distance_metric::DistanceMetric`: Distance metric for localization
- `adaptive_scaling::Bool`: Enable adaptive radius scaling
- `ensemble_based_adaptation::Bool`: Use ensemble spread for adaptation
- `multi_variable_scaling::Dict{Tuple{Symbol,Symbol}, T}`: Cross-variable scaling factors
- `vertical_localization_scales::Vector{T}`: Vertical localization scales by level
- `horizontal_anisotropy_ratio::T`: Horizontal anisotropy ratio (Lx/Ly)
"""
struct LocalizationConfiguration{T<:AbstractFloat}
    base_radius::T
    distance_metric::DistanceMetric
    adaptive_scaling::Bool
    ensemble_based_adaptation::Bool
    multi_variable_scaling::Dict{Tuple{Symbol,Symbol}, T}
    vertical_localization_scales::Vector{T}
    horizontal_anisotropy_ratio::T
    
    function LocalizationConfiguration{T}(;
        base_radius::T = T(500_000),                    # 500 km base radius
        distance_metric::DistanceMetric = EUCLIDEAN,
        adaptive_scaling::Bool = true,
        ensemble_based_adaptation::Bool = false,
        multi_variable_scaling::Dict{Tuple{Symbol,Symbol}, T} = Dict{Tuple{Symbol,Symbol}, T}(),
        vertical_localization_scales::Vector{T} = T[],
        horizontal_anisotropy_ratio::T = T(1)
    ) where T
        new{T}(base_radius, distance_metric, adaptive_scaling, ensemble_based_adaptation,
               multi_variable_scaling, vertical_localization_scales, horizontal_anisotropy_ratio)
    end
end

LocalizationConfiguration(args...; kwargs...) = LocalizationConfiguration{Float64}(args...; kwargs...)

"""
    LocalizationOperator{T<:AbstractFloat}

Sparse localization operator for efficient covariance localization.

# Fields
- `localization_matrix::SparseMatrixCSC{T,Int}`: Sparse localization matrix
- `localization_function::AbstractLocalizationFunction{T}`: Localization function
- `config::LocalizationConfiguration{T}`: Configuration parameters
- `grid_coordinates::Matrix{T}`: Grid point coordinates (n_points × n_dims)
- `distance_cache::Dict{Tuple{Int,Int}, T}`: Cached distance computations
"""
struct LocalizationOperator{T<:AbstractFloat}
    localization_matrix::SparseMatrixCSC{T,Int}
    localization_function::AbstractLocalizationFunction{T}
    config::LocalizationConfiguration{T}
    grid_coordinates::Matrix{T}
    distance_cache::Dict{Tuple{Int,Int}, T}
    
    function LocalizationOperator{T}(localization_function::AbstractLocalizationFunction{T},
                                    config::LocalizationConfiguration{T},
                                    grid_coordinates::Matrix{T}) where T
        
        n_points = size(grid_coordinates, 1)
        
        # Build sparse localization matrix
        loc_matrix, dist_cache = build_localization_matrix(localization_function, config, 
                                                          grid_coordinates)
        
        new{T}(loc_matrix, localization_function, config, grid_coordinates, dist_cache)
    end
end

"""
    AdaptiveLocalization{T<:AbstractFloat}

Adaptive localization with flow-dependent and ensemble-spread-dependent radius.

# Fields
- `base_operator::LocalizationOperator{T}`: Base localization operator
- `adaptation_factors::Matrix{T}`: Spatially varying adaptation factors
- `ensemble_spread::Matrix{T}`: Local ensemble spread estimates
- `flow_dependent_scaling::Matrix{T}`: Flow-dependent radius scaling
"""
mutable struct AdaptiveLocalization{T<:AbstractFloat}
    base_operator::LocalizationOperator{T}
    adaptation_factors::Matrix{T}
    ensemble_spread::Union{Matrix{T}, Nothing}
    flow_dependent_scaling::Union{Matrix{T}, Nothing}
    
    function AdaptiveLocalization{T}(base_operator::LocalizationOperator{T}) where T
        n_points = size(base_operator.grid_coordinates, 1)
        adaptation_factors = ones(T, n_points, n_points)
        
        new{T}(base_operator, adaptation_factors, nothing, nothing)
    end
end

"""
    MultiScaleLocalization{T<:AbstractFloat}

Multi-scale localization for hierarchical covariance structures.

# Fields
- `scale_operators::Vector{LocalizationOperator{T}}`: Operators for each scale
- `scale_weights::Vector{T}`: Weights for combining scales
- `scale_radii::Vector{T}`: Localization radii for each scale
"""
struct MultiScaleLocalization{T<:AbstractFloat}
    scale_operators::Vector{LocalizationOperator{T}}
    scale_weights::Vector{T}
    scale_radii::Vector{T}
    
    function MultiScaleLocalization{T}(base_function::AbstractLocalizationFunction{T},
                                      config::LocalizationConfiguration{T},
                                      grid_coordinates::Matrix{T},
                                      scale_radii::Vector{T}) where T
        
        # Create localization function and operator for each scale
        operators = LocalizationOperator{T}[]
        for radius in scale_radii
            # Create scaled localization function
            if base_function isa GaspariCohnFunction
                scaled_function = GaspariCohnFunction{T}(radius, base_function.cutoff_ratio)
            elseif base_function isa ExponentialFunction
                scaled_function = ExponentialFunction{T}(radius, base_function.decay_rate, 
                                                       base_function.cutoff_threshold)
            else
                error("Unsupported localization function type")
            end
            
            # Create operator for this scale
            operator = LocalizationOperator{T}(scaled_function, config, grid_coordinates)
            push!(operators, operator)
        end
        
        # Equal weights by default - could be optimized
        weights = ones(T, length(scale_radii)) ./ length(scale_radii)
        
        new{T}(operators, weights, scale_radii)
    end
end

"""
    (func::GaspariCohnFunction{T})(distance::T) where T

Evaluate Gaspari-Cohn localization function at given distance.
"""
function (func::GaspariCohnFunction{T})(distance::T) where T
    r = distance / func.cutoff_radius
    
    if r <= 1
        return -T(0.25)*r^5 + T(0.5)*r^4 + T(5/8)*r^3 - T(5/3)*r^2 + T(1)
    elseif r <= func.cutoff_ratio
        return T(1/12)*r^5 - T(0.5)*r^4 + T(5/8)*r^3 + T(5/3)*r^2 - 
               T(5)*r + T(4) - T(2/3)*(func.cutoff_radius/distance)
    else
        return T(0)
    end
end

"""
    (func::ExponentialFunction{T})(distance::T) where T

Evaluate exponential localization function at given distance.
"""
function (func::ExponentialFunction{T})(distance::T) where T
    normalized_distance = distance / func.length_scale
    correlation = exp(-func.decay_rate * normalized_distance^2 / 2)
    
    return correlation > func.cutoff_threshold ? correlation : T(0)
end

"""
    compute_distance(coord1::AbstractVector{T}, coord2::AbstractVector{T}, 
                    metric::DistanceMetric) where T

Compute distance between two coordinate points using specified metric.
"""
function compute_distance(coord1::AbstractVector{T}, coord2::AbstractVector{T}, 
                         metric::DistanceMetric) where T
    if metric == EUCLIDEAN
        return sqrt(sum((coord1 - coord2).^2))
    elseif metric == GREAT_CIRCLE
        # Assume coordinates are [lon, lat] in radians
        lat1, lon1 = coord1[2], coord1[1]
        lat2, lon2 = coord2[2], coord2[1]
        
        # Haversine formula
        dlat = lat2 - lat1
        dlon = lon2 - lon1
        a = sin(dlat/2)^2 + cos(lat1) * cos(lat2) * sin(dlon/2)^2
        c = 2 * asin(sqrt(a))
        
        # Earth radius in meters
        R = T(6.371e6)
        return R * c
    elseif metric == MANHATTAN
        return sum(abs.(coord1 - coord2))
    elseif metric == MAXIMUM
        return maximum(abs.(coord1 - coord2))
    else
        error("Unsupported distance metric: $metric")
    end
end

"""
    build_localization_matrix(func::AbstractLocalizationFunction{T},
                             config::LocalizationConfiguration{T},
                             grid_coordinates::Matrix{T}) where T

Build sparse localization matrix from localization function and grid coordinates.
"""
function build_localization_matrix(func::AbstractLocalizationFunction{T},
                                  config::LocalizationConfiguration{T},
                                  grid_coordinates::Matrix{T}) where T
    n_points = size(grid_coordinates, 1)
    
    # Pre-allocate arrays for sparse matrix construction
    I_indices = Int[]
    J_indices = Int[]
    values = T[]
    distance_cache = Dict{Tuple{Int,Int}, T}()
    
    # Determine maximum localization radius for efficiency
    if func isa GaspariCohnFunction
        max_radius = func.cutoff_radius * func.cutoff_ratio
    elseif func isa ExponentialFunction
        max_radius = func.length_scale * sqrt(-2 * log(func.cutoff_threshold))
    else
        max_radius = config.base_radius * T(3)  # Default: 3 times base radius
    end
    
    # Build localization matrix entries
    for i in 1:n_points
        for j in i:n_points  # Exploit symmetry
            # Compute distance
            coord_i = view(grid_coordinates, i, :)
            coord_j = view(grid_coordinates, j, :)
            distance = compute_distance(coord_i, coord_j, config.distance_metric)
            
            # Cache distance computation
            distance_cache[(i,j)] = distance
            if i != j
                distance_cache[(j,i)] = distance
            end
            
            # Skip if distance exceeds localization radius
            if distance > max_radius
                continue
            end
            
            # Apply horizontal anisotropy if specified
            if config.horizontal_anisotropy_ratio != 1
                # Simple anisotropy scaling (could be made more sophisticated)
                if length(coord_i) >= 2
                    dx = coord_j[1] - coord_i[1]
                    dy = coord_j[2] - coord_i[2]
                    anisotropic_distance = sqrt((dx/config.horizontal_anisotropy_ratio)^2 + dy^2)
                    distance = anisotropic_distance
                end
            end
            
            # Evaluate localization function
            loc_value = func(distance)
            
            if loc_value > 1e-12  # Numerical cutoff
                push!(I_indices, i)
                push!(J_indices, j)  
                push!(values, loc_value)
                
                # Add symmetric entry if not diagonal
                if i != j
                    push!(I_indices, j)
                    push!(J_indices, i)
                    push!(values, loc_value)
                end
            end
        end
    end
    
    # Construct sparse matrix
    loc_matrix = sparse(I_indices, J_indices, values, n_points, n_points)
    
    return loc_matrix, distance_cache
end

"""
    apply_localization(operator::LocalizationOperator{T}, 
                      covariance_matrix::AbstractMatrix{T}) where T

Apply localization to covariance matrix using Schur product.
"""
function apply_localization(operator::LocalizationOperator{T}, 
                           covariance_matrix::AbstractMatrix{T}) where T
    # Schur (element-wise) product
    return operator.localization_matrix .* covariance_matrix
end

"""
    apply_localization(operator::LocalizationOperator{T}, 
                      ensemble_perturbations::Matrix{T}) where T

Apply localization to ensemble covariance implicitly through perturbations.
"""
function apply_localization(operator::LocalizationOperator{T}, 
                           ensemble_perturbations::Matrix{T}) where T
    n_grid, n_ens = size(ensemble_perturbations)
    
    # Apply localization through matrix square root approach
    # This is more efficient than forming full covariance matrix
    
    # Compute localized ensemble perturbations
    # Using Cholesky-like decomposition of localization matrix
    localized_perturbations = similar(ensemble_perturbations)
    
    # Simple approach: apply localization to each ensemble member
    # More sophisticated methods could use matrix square root
    for ens in 1:n_ens
        perturbation = ensemble_perturbations[:, ens]
        
        # Apply spatial localization (simplified)
        localized_perturbation = apply_spatial_localization(operator, perturbation)
        localized_perturbations[:, ens] = localized_perturbation
    end
    
    return localized_perturbations
end

"""
    apply_spatial_localization(operator::LocalizationOperator{T}, field::Vector{T}) where T

Apply spatial localization to a field vector.
"""
function apply_spatial_localization(operator::LocalizationOperator{T}, field::Vector{T}) where T
    # This is a simplified approach - in practice would use more sophisticated methods
    n_points = length(field)
    localized_field = zeros(T, n_points)
    
    # Apply localization through weighted averaging
    rows = rowvals(operator.localization_matrix)
    vals = nonzeros(operator.localization_matrix)
    
    for i in 1:n_points
        sum_weights = T(0)
        weighted_sum = T(0)
        
        for idx in nzrange(operator.localization_matrix, i)
            j = rows[idx]
            weight = vals[idx]
            
            weighted_sum += weight * field[j]
            sum_weights += weight
        end
        
        if sum_weights > 0
            localized_field[i] = weighted_sum / sum_weights
        else
            localized_field[i] = field[i]
        end
    end
    
    return localized_field
end

"""
    update_adaptive_factors!(adaptive::AdaptiveLocalization{T}, 
                            ensemble_perturbations::Matrix{T}) where T

Update adaptive localization factors based on ensemble spread.
"""
function update_adaptive_factors!(adaptive::AdaptiveLocalization{T}, 
                                 ensemble_perturbations::Matrix{T}) where T
    if !adaptive.base_operator.config.adaptive_scaling
        return
    end
    
    n_grid, n_ens = size(ensemble_perturbations)
    
    # Compute local ensemble spread
    ensemble_spread = std(ensemble_perturbations, dims=2)[:, 1]
    adaptive.ensemble_spread = reshape(ensemble_spread, size(adaptive.adaptation_factors, 1), 1)
    
    # Compute adaptation factors based on spread
    mean_spread = mean(ensemble_spread)
    
    for i in 1:n_grid
        # Increase localization radius where spread is large
        # Decrease where spread is small
        spread_ratio = ensemble_spread[i] / mean_spread
        adaptation_factor = T(0.5) + T(0.5) * tanh(T(2) * (spread_ratio - T(1)))
        
        adaptive.adaptation_factors[i, :] .= adaptation_factor
        adaptive.adaptation_factors[:, i] .= adaptation_factor
    end
end

"""
    apply_localization(adaptive::AdaptiveLocalization{T}, 
                      covariance_matrix::AbstractMatrix{T}) where T

Apply adaptive localization with spatially varying localization radius.
"""
function apply_localization(adaptive::AdaptiveLocalization{T}, 
                           covariance_matrix::AbstractMatrix{T}) where T
    # Apply base localization
    base_localized = apply_localization(adaptive.base_operator, covariance_matrix)
    
    # Apply adaptive scaling
    return base_localized .* adaptive.adaptation_factors
end

"""
    apply_localization(multi::MultiScaleLocalization{T}, 
                      covariance_matrix::AbstractMatrix{T}) where T

Apply multi-scale localization combining different length scales.
"""
function apply_localization(multi::MultiScaleLocalization{T}, 
                           covariance_matrix::AbstractMatrix{T}) where T
    n_scales = length(multi.scale_operators)
    result = zeros(T, size(covariance_matrix))
    
    # Combine localization across scales
    for (i, operator) in enumerate(multi.scale_operators)
        weight = multi.scale_weights[i]
        scale_localized = apply_localization(operator, covariance_matrix)
        result .+= weight .* scale_localized
    end
    
    return result
end

"""
    create_localization_operator(func::AbstractLocalizationFunction{T},
                                config::LocalizationConfiguration{T},
                                grid_coordinates::Matrix{T}) where T

Factory function to create localization operator.
"""
function create_localization_operator(func::AbstractLocalizationFunction{T},
                                     config::LocalizationConfiguration{T},
                                     grid_coordinates::Matrix{T}) where T
    return LocalizationOperator{T}(func, config, grid_coordinates)
end

"""
    compute_localization_matrix(func::AbstractLocalizationFunction{T},
                               coordinates::Matrix{T};
                               distance_metric::DistanceMetric = EUCLIDEAN) where T

Convenience function to compute localization matrix directly.
"""
function compute_localization_matrix(func::AbstractLocalizationFunction{T},
                                   coordinates::Matrix{T};
                                   distance_metric::DistanceMetric = EUCLIDEAN) where T
    config = LocalizationConfiguration{T}(distance_metric=distance_metric)
    operator = create_localization_operator(func, config, coordinates)
    return operator.localization_matrix
end

"""
    compute_localized_covariance(raw_covariance::Matrix{T}, 
                                localization_function::AbstractLocalizationFunction{T},
                                grid_coordinates::Matrix{T}) where T

Compute localized covariance matrix in one step.
"""
function compute_localized_covariance(raw_covariance::Matrix{T}, 
                                     localization_function::AbstractLocalizationFunction{T},
                                     grid_coordinates::Matrix{T}) where T
    config = LocalizationConfiguration{T}()
    operator = create_localization_operator(localization_function, config, grid_coordinates)
    return apply_localization(operator, raw_covariance)
end

"""
    validate_localization_function(func::AbstractLocalizationFunction{T}, 
                                  test_distances::Vector{T}) where T

Validate localization function properties and behavior.
"""
function validate_localization_function(func::AbstractLocalizationFunction{T}, 
                                       test_distances::Vector{T}) where T
    n_test = length(test_distances)
    function_values = T[func(d) for d in test_distances]
    
    # Check basic properties
    monotonic_decreasing = all(i -> function_values[i] >= function_values[i+1], 1:n_test-1)
    positive_semidefinite = all(>=)(function_values, 0)
    unit_at_origin = abs(func(T(0)) - T(1)) < 1e-12
    
    # Compute effective support radius
    cutoff_idx = findfirst(<(1e-6), function_values)
    effective_support = cutoff_idx !== nothing ? test_distances[cutoff_idx] : test_distances[end]
    
    return (values=function_values, 
            monotonic=monotonic_decreasing,
            positive=positive_semidefinite,
            unit_origin=unit_at_origin,
            support_radius=effective_support)
end

end # module CovarianceLocalization